import numpy as np
import math

class LinkBudgetCalculator:
    """链路预算计算器"""
    
    def __init__(self):
        # 基本参数
        self.tx_power = 30  # dBm
        self.frequency = 2.4  # GHz
        
        # 天线参数
        self.tx_antenna_gain = 15  # dBi
        self.rx_antenna_gain = 35  # dBi
        
        # 系统参数
        self.bandwidth = 1e6  # Hz
        self.noise_figure = 2  # dB
        self.required_eb_no = 10  # dB
        
        # 大气损耗
        self.atmospheric_loss = 0.5  # dB
        self.rain_loss = 1.0  # dB
        self.polarization_loss = 0.5  # dB
        
    def set_parameters(self, tx_power, frequency):
        """设置通信参数"""
        self.tx_power = tx_power
        self.frequency = frequency
        
    def calculate_link_budget(self, satellite_position):
        """计算链路预算"""
        # 计算距离（假设地面站在赤道）
        ground_station = {
            'latitude': 0,
            'longitude': 0,
            'altitude': 0
        }
        
        distance = self._calculate_distance(satellite_position, ground_station)
        
        # 计算自由空间损耗
        fspl = self._calculate_fspl(distance)
        
        # 计算接收功率
        rx_power = self.tx_power + self.tx_antenna_gain + self.rx_antenna_gain - fspl - \
                   self.atmospheric_loss - self.rain_loss - self.polarization_loss
        
        # 计算噪声功率
        noise_power = self._calculate_noise_power()
        
        # 计算信噪比
        snr = rx_power - noise_power
        
        # 计算数据速率
        max_data_rate = self._calculate_max_data_rate(snr)
        
        # 计算链路余量
        link_margin = snr - self.required_eb_no
        
        return {
            'distance': distance,
            'fspl': fspl,
            'rx_power': rx_power,
            'noise_power': noise_power,
            'snr': snr,
            'max_data_rate': max_data_rate,
            'link_margin': link_margin,
            'elevation_angle': self._calculate_elevation_angle(satellite_position, ground_station)
        }
    
    def _calculate_distance(self, sat_pos, ground_station):
        """计算卫星到地面站的距离"""
        # 简化的距离计算
        earth_radius = 6371  # km
        satellite_altitude = sat_pos['altitude']
        
        # 计算地面站与卫星的直线距离
        lat1 = math.radians(ground_station['latitude'])
        lon1 = math.radians(ground_station['longitude'])
        lat2 = math.radians(sat_pos['latitude'])
        lon2 = math.radians(sat_pos['longitude'])
        
        # 使用球面距离公式
        dlat = lat2 - lat1
        dlon = lon2 - lon1
        
        a = math.sin(dlat/2)**2 + math.cos(lat1) * math.cos(lat2) * math.sin(dlon/2)**2
        c = 2 * math.atan2(math.sqrt(a), math.sqrt(1-a))
        
        # 计算斜距
        ground_distance = earth_radius * c
        distance = math.sqrt(ground_distance**2 + satellite_altitude**2)
        
        return distance
    
    def _calculate_fspl(self, distance):
        """计算自由空间路径损耗"""
        # 转换为米
        distance_m = distance * 1000
        frequency_hz = self.frequency * 1e9
        
        # 自由空间路径损耗公式
        fspl = 20 * np.log10(distance_m) + 20 * np.log10(frequency_hz) - 147.55
        
        return fspl
    
    def _calculate_noise_power(self):
        """计算噪声功率"""
        # 玻尔兹曼常数
        k = 1.38e-23  # J/K
        T = 290  # K (标准温度)
        
        # 噪声功率 (dBm)
        noise_power_dbm = 10 * np.log10(k * T * self.bandwidth * 1000)
        noise_power_dbm += self.noise_figure
        
        return noise_power_dbm
    
    def _calculate_max_data_rate(self, snr):
        """计算最大数据速率 (Shannon定理)"""
        # 转换为线性值
        snr_linear = 10**(snr/10)
        
        # Shannon容量
        capacity = self.bandwidth * np.log2(1 + snr_linear)
        
        return capacity
    
    def _calculate_elevation_angle(self, sat_pos, ground_station):
        """计算仰角"""
        # 计算卫星到地面站的向量
        earth_radius = 6371  # km
        
        # 计算地面站坐标
        gs_x = earth_radius * np.cos(np.radians(ground_station['latitude'])) * np.cos(np.radians(ground_station['longitude']))
        gs_y = earth_radius * np.cos(np.radians(ground_station['latitude'])) * np.sin(np.radians(ground_station['longitude']))
        gs_z = earth_radius * np.sin(np.radians(ground_station['latitude']))
        
        # 计算卫星坐标
        sat_x = (earth_radius + sat_pos['altitude']) * np.cos(np.radians(sat_pos['latitude'])) * np.cos(np.radians(sat_pos['longitude']))
        sat_y = (earth_radius + sat_pos['altitude']) * np.cos(np.radians(sat_pos['latitude'])) * np.sin(np.radians(sat_pos['longitude']))
        sat_z = (earth_radius + sat_pos['altitude']) * np.sin(np.radians(sat_pos['latitude']))
        
        # 计算向量
        gs_to_sat = np.array([sat_x - gs_x, sat_y - gs_y, sat_z - gs_z])
        gs_to_center = np.array([gs_x, gs_y, gs_z])
        
        # 计算仰角
        cos_elevation = np.dot(gs_to_sat, gs_to_center) / (np.linalg.norm(gs_to_sat) * np.linalg.norm(gs_to_center))
        elevation = 90 - np.degrees(np.arccos(np.clip(cos_elevation, -1, 1)))
        
        return max(elevation, 0)